#carregamento de todas as imagens

import cv2
import numpy as np
import matplotlib.pyplot as plt



def exibir_imagens_lado_a_lado(imagens, titulos, figsize=(20, 5), cmap=None):

    # Criar subplots com o número de imagens fornecido
    fig, axs = plt.subplots(1, len(imagens), figsize=figsize)

    # Garantir que axs seja uma lista (caso haja apenas uma imagem)
    if len(imagens) == 1:
        axs = [axs]

    # Exibir cada imagem no subplot correspondente
    for ax, img, titulo in zip(axs, imagens, titulos):
        # Exibir a imagem
        ax.imshow(img, cmap=cmap)
        # Configurar o título e desativar os eixos
        ax.set_title(titulo)
        ax.axis('off')

    plt.tight_layout()
    plt.show()


def carregar_imagens(caminho_base, quantidade, formato="RGB"):
    imagens = []

    for i in range(1, quantidade + 1):
        caminho = f"{caminho_base}{i}.jpg"
        imagem = cv2.imread(caminho)

        if imagem is None:
            print(f"Erro ao carregar imagem {caminho}")
            continue

        # Converter o formato de cor, se necessário
        if formato.upper() == "RGB":
            imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2RGB)
        elif formato.upper() == "GRAY":
            imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)

        imagens.append(imagem)

    return imagens

canetas = carregar_imagens("canetas/caneta", 10)
borrachas = carregar_imagens("borracha/borracha", 10)
livros = carregar_imagens("livros/livro", 10)

caneta_titulos = [f"Caneta{idx}" for idx in range(1, len(canetas) + 1)]
exibir_imagens_lado_a_lado(canetas, caneta_titulos)

borracha_titulos = [f"Borracha {idx}" for idx in range(1, len(borrachas) + 1)]
exibir_imagens_lado_a_lado(borrachas, borracha_titulos)

livro_titulos = [f"Livro {idx}" for idx in range(1, len(livros) + 1)]
exibir_imagens_lado_a_lado(livros, livro_titulos)

#data augmentation das canetas

# Diretórios para salvar as imagens
C_espelhamento = "canetas_espelhamento/"
C_zoom = "canetas_zoom/"
C_rotacao = "canetas_rotacao/"


# Função para aplicar espelhamento horizontal
def flip_image(image):
    return cv2.flip(image, 1)

# Função para aplicar zoom
def zoom_image(image, zoom_factor=1.2):
    height, width = image.shape[:2]
    center_x, center_y = width // 2, height // 2
    new_width, new_height = int(width / zoom_factor), int(height / zoom_factor)
    start_x = center_x - new_width // 2
    start_y = center_y - new_height // 2
    cropped = image[start_y:start_y + new_height, start_x:start_x + new_width]
    return cv2.resize(cropped, (width, height))

# Função para aplicar rotação
def rotate_image(image, angle):
    height, width = image.shape[:2]
    center = (width // 2, height // 2)
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    return cv2.warpAffine(image, rotation_matrix, (width, height),borderValue=(255, 255, 255))

# Listas para armazenar imagens transformadas
canetasEspelhamento = []
canetasZoom = []
canetasRotacao = []

# Carregar e processar imagens
for i in range(len(canetas)):

    # Aplicar espelhamento
    imagem_flip = flip_image(canetas[i])
    canetasEspelhamento.append(imagem_flip)
    caminho_flip = f"{C_espelhamento}caneta_espelhamento{i+1}.jpg"
    cv2.imwrite(caminho_flip, cv2.cvtColor(imagem_flip, cv2.COLOR_RGB2BGR))
    
    # Aplicar zoom
    imagem_zoom = zoom_image(canetas[i])
    canetasZoom.append(imagem_zoom)
    caminho_zoom = f"{C_zoom}caneta_zoom{i+1}.jpg"
    cv2.imwrite(caminho_zoom, cv2.cvtColor(imagem_zoom, cv2.COLOR_RGB2BGR))
    # Aplicar rotação
    imagem_rotacao = rotate_image(canetas[i], angle=15)  # Exemplo: rotação de 15 graus
    canetasRotacao.append(imagem_rotacao)
    caminho_rotacao = f"{C_rotacao}caneta_rotacao{i+1}.jpg"
    cv2.imwrite(caminho_rotacao, cv2.cvtColor(imagem_rotacao, cv2.COLOR_RGB2BGR))

canetas_espelhamento = [f"{idx} espelhamento" for idx in range(1, len(canetasEspelhamento) + 1)]
exibir_imagens_lado_a_lado(canetasEspelhamento, canetas_espelhamento)

canetas_zoom = [f"{idx} zoom" for idx in range(1, len(canetasZoom) + 1)]
exibir_imagens_lado_a_lado(canetasZoom, canetas_zoom)

canetas_rotacao = [f"{idx} rotação" for idx in range(1, len(canetasRotacao) + 1)]
exibir_imagens_lado_a_lado(canetasRotacao, canetas_rotacao )

#data augmentation das borrachas

# Diretórios para salvar as imagens
B_espelhamento = "borrachas_espelhamento/"
B_zoom = "borrachas_zoom/"
B_rotacao = "borrachas_rotacao/"

# Listas para armazenar imagens transformadas
borrachasEspelhamento = []
borrachasZoom = []
borrachasRotacao = []

# Carregar e processar imagens
for i in range(len(borrachas)):

    # Aplicar espelhamento
    imagem_flip = flip_image(borrachas[i])
    borrachasEspelhamento.append(imagem_flip)
    caminho_flip = f"{B_espelhamento}borracha_espelhamento{i+1}.jpg"
    cv2.imwrite(caminho_flip, cv2.cvtColor(imagem_flip, cv2.COLOR_RGB2BGR))
    
    # Aplicar zoom
    imagem_zoom = zoom_image(borrachas[i])
    borrachasZoom.append(imagem_zoom)
    caminho_zoom = f"{B_zoom}borracha_zoom{i+1}.jpg"
    cv2.imwrite(caminho_zoom, cv2.cvtColor(imagem_zoom, cv2.COLOR_RGB2BGR))
    
    # Aplicar rotação
    imagem_rotacao = rotate_image(borrachas[i], angle=15)  # Exemplo: rotação de 15 graus
    borrachasRotacao.append(imagem_rotacao)
    caminho_rotacao = f"{B_rotacao}borracha_rotacao{i+1}.jpg"
    cv2.imwrite(caminho_rotacao, cv2.cvtColor(imagem_rotacao, cv2.COLOR_RGB2BGR))

borrachas_espelhamento = [f"{idx} espelhamento" for idx in range(1, len(borrachasEspelhamento) + 1)]
exibir_imagens_lado_a_lado(borrachasEspelhamento, borrachas_espelhamento)

borrachas_zoom = [f"{idx} zoom" for idx in range(1, len(borrachasZoom) + 1)]
exibir_imagens_lado_a_lado(borrachasZoom, borrachas_zoom)

borrachas_rotacao = [f"{idx} rotação" for idx in range(1, len(borrachasRotacao) + 1)]
exibir_imagens_lado_a_lado(borrachasRotacao, borrachas_rotacao )

#data augmentation dos livros

# Diretórios para salvar as imagens
L_espelhamento = "livros_espelhamento/"
L_zoom = "livros_zoom/"
L_rotacao = "livros_rotacao/"


# Listas para armazenar imagens transformadas
livrosEspelhamento = []
livrosZoom = []
livrosRotacao = []

# Carregar e processar imagens
for i in range(len(livros)):

    # Aplicar espelhamento
    imagem_flip = flip_image(livros[i])
    livrosEspelhamento.append(imagem_flip)
    caminho_flip = f"{L_espelhamento}livro_espelhamento{i+1}.jpg"
    cv2.imwrite(caminho_flip, cv2.cvtColor(imagem_flip, cv2.COLOR_RGB2BGR))

    # Aplicar zoom
    imagem_zoom = zoom_image(livros[i])
    livrosZoom.append(imagem_zoom)
    caminho_zoom = f"{L_zoom}livro_zoom{i+1}.jpg"
    cv2.imwrite(caminho_zoom, cv2.cvtColor(imagem_zoom, cv2.COLOR_RGB2BGR))
    
    # Aplicar rotação
    imagem_rotacao = rotate_image(livros[i], angle=15)  # Exemplo: rotação de 15 graus
    livrosRotacao.append(imagem_rotacao)
    caminho_rotacao = f"{L_rotacao}livro_rotacao{i+1}.jpg"
    cv2.imwrite(caminho_rotacao, cv2.cvtColor(imagem_rotacao, cv2.COLOR_RGB2BGR))

livros_espelhamento = [f"{idx} espelhamento" for idx in range(1, len(livrosEspelhamento) + 1)]
exibir_imagens_lado_a_lado(livrosEspelhamento, livros_espelhamento)

livros_zoom = [f"{idx} zoom" for idx in range(1, len(livrosZoom) + 1)]
exibir_imagens_lado_a_lado(livrosZoom, livros_zoom)

livros_rotacao = [f"{idx} rotação" for idx in range(1, len(livrosRotacao) + 1)]
exibir_imagens_lado_a_lado(livrosRotacao, livros_rotacao )

#segmentação das imagens originais
livrosSegmentados = []
canetasSegmentadas = []
borrachasSegmentadas = []

#limiarização da imagens dos livros
for i in range(1, 11):
    caminho = f'livros/livro{i}.jpg'
    imagem = cv2.imread(caminho)

    if imagem is None:
        print(f'erro ao carregar imagem {caminho}')
        continue
    # Convertendo para escala de cinza
    imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)
    # Aplicando filtro Gaussiano
    imagem_blur = cv2.GaussianBlur(imagem, (5, 5), 0)
    # Aplicando limiarização
    _, thresh = cv2.threshold(imagem_blur, 120, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((25, 25), np.uint8)
    dilatacao = cv2.dilate(thresh, kernel, iterations=30)
    thresh = cv2.erode(dilatacao, kernel, iterations=40) 
    thresh = cv2.dilate(thresh, kernel, iterations=10) 
    livrosSegmentados.append(thresh)  # Armazenar a imagem processada

#limiarização da imagens das borrachas
for i in range(1, 11):
    caminho = f'borracha/borracha{i}.jpg'
    imagem = cv2.imread(caminho)

    if imagem is None:
        print(f'erro ao carregar imagem {caminho}')
        continue
    # Convertendo para escala de cinza
    imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)
    # Aplicando filtro Gaussiano
    imagem_blur = cv2.GaussianBlur(imagem, (5, 5), 0)
    # Aplicando limiarização
    _, thresh = cv2.threshold(imagem_blur, 125, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((15, 15), np.uint8)
    dilatacao = cv2.dilate(thresh, kernel, iterations=5)
    thresh = cv2.erode(dilatacao, kernel, iterations=6) 
    borrachasSegmentadas.append(thresh)  # Armazenar a imagem processada

#limiarização da imagens das canetas
for i in range(1, 11):  # De 1 a 10
    caminho = f'canetas/caneta{i}.jpg'  
    imagem = cv2.imread(caminho)
    
    if imagem is None:
        print(f"Erro ao carregar a imagem {caminho}")
        continue

    imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)
    imagem_blur = cv2.GaussianBlur(imagem, (5, 5), 0)
    _, thresh = cv2.threshold(imagem_blur, 125, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((15, 15), np.uint8)
    dilatacao = cv2.dilate(thresh, kernel, iterations=5)
    thresh = cv2.erode(dilatacao, kernel, iterations=6)
    canetasSegmentadas.append(thresh)  


caneta_segmentada = [f"limiarizada{idx}" for idx in range(1, len(canetasSegmentadas) + 1)]
exibir_imagens_lado_a_lado(canetasSegmentadas, caneta_segmentada, cmap='gray')

borracha_segmentada = [f"limiarizada{idx}" for idx in range(1, len(borrachasSegmentadas) + 1)]
exibir_imagens_lado_a_lado(borrachasSegmentadas, borracha_segmentada, cmap='gray')

livro_segmentada = [f"limiarizada{idx}" for idx in range(1, len(livrosSegmentados) + 1)]
exibir_imagens_lado_a_lado(livrosSegmentados, livro_segmentada, cmap='gray')

#segmentação das canetas em data augmentation

canetasEsp = carregar_imagens("canetas_espelhamento/caneta_espelhamento", 10)
canetasZooms = carregar_imagens("canetas_zoom/caneta_zoom", 10)
canetasRota = carregar_imagens("canetas_rotacao/caneta_rotacao", 10)

canetasSegEspelhamento = []
canetasSegZoom = []
canetasSegRotacao = []

#limiarização da imagens das canetas
for i in range(1, 11):  # De 1 a 10
    caminho = f'canetas_espelhamento/caneta_espelhamento{i}.jpg'  
    imagem = cv2.imread(caminho)
    
    if imagem is None:
        print(f"Erro ao carregar a imagem {caminho}")
        continue

    imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)
    imagem_blur = cv2.GaussianBlur(imagem, (5, 5), 0)
    _, thresh = cv2.threshold(imagem_blur, 125, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((15, 15), np.uint8)
    dilatacao = cv2.dilate(thresh, kernel, iterations=5)
    thresh = cv2.erode(dilatacao, kernel, iterations=6)
    canetasSegEspelhamento.append(thresh)  

#limiarização da imagens das canetas
for i in range(1, 11):  # De 1 a 10
    caminho = f'canetas_zoom/caneta_zoom{i}.jpg'  
    imagem = cv2.imread(caminho)
    
    if imagem is None:
        print(f"Erro ao carregar a imagem {caminho}")
        continue

    imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)
    imagem_blur = cv2.GaussianBlur(imagem, (5, 5), 0)
    _, thresh = cv2.threshold(imagem_blur, 125, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((15, 15), np.uint8)
    dilatacao = cv2.dilate(thresh, kernel, iterations=5)
    thresh = cv2.erode(dilatacao, kernel, iterations=6)
    canetasSegZoom.append(thresh)  

#limiarização da imagens das canetas
for i in range(1, 11):  # De 1 a 10
    caminho = f'canetas_rotacao/caneta_rotacao{i}.jpg'  
    imagem = cv2.imread(caminho)
    
    if imagem is None:
        print(f"Erro ao carregar a imagem {caminho}")
        continue

    imagem = cv2.cvtColor(imagem, cv2.COLOR_BGR2GRAY)
    imagem_blur = cv2.GaussianBlur(imagem, (5, 5), 0)
    _, thresh = cv2.threshold(imagem_blur, 125, 255, cv2.THRESH_BINARY_INV)
    kernel = np.ones((15, 15), np.uint8)
    dilatacao = cv2.dilate(thresh, kernel, iterations=5)
    thresh = cv2.erode(dilatacao, kernel, iterations=6)
    canetasSegRotacao.append(thresh)  


caneta_segEsp = [f"espelhamento {idx}" for idx in range(1, len(canetasSegEspelhamento) + 1)]
exibir_imagens_lado_a_lado(canetasSegEspelhamento, caneta_segEsp, cmap='gray')

caneta_segZoom = [f"zoom {idx}" for idx in range(1, len(canetasSegZoom) + 1)]
exibir_imagens_lado_a_lado(canetasSegZoom, caneta_segZoom, cmap='gray')

caneta_segRota = [f"rotação {idx}" for idx in range(1, len(canetasSegRotacao) + 1)]
exibir_imagens_lado_a_lado(canetasSegRotacao, caneta_segRota, cmap='gray')